
#include "asinh_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling
{
    const uint32_t BLOCK_SIZE = 32;
    const uint32_t BUFFER_NUM = 2;
    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {
        AsinhTilingData tiling;

        uint64_t ubSize;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);

        // 获取输入数据的shapesize    shape(1, 660), shapesize = 660
        uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
        // 获取输入数据的类型 如 half 2
        uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
        // 数据的大小 1320
        uint32_t inputLength = inputNum * inputBytes;

        // 使用的变量数，先暂设5，根据类型进行修改
        uint32_t ubDataNumber = 4;

        ge::DataType dataType = context->GetInputTensor(0)->GetDataType();
        switch (dataType)
        {
        case ge::DT_FLOAT:
            ubDataNumber = 3;
            context->SetTilingKey(1);
            break;
        case ge::DT_FLOAT16:
            ubDataNumber = 6;
            context->SetTilingKey(2);
            break;
        default:
            break;
        }

        // 单核一次计算的block数 8
        uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;

        // 单次搬运的数据量，一次处理多少个数 8 * 32  / 2 = 128
        uint32_t tileDataNum = tileBlockNum * BLOCK_SIZE / inputBytes;

        // 输入数据 32B 向上对齐 1344
        uint32_t inputLengthAlgin32 = ((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;

        // 一个核需要处理多少个block，42
        uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE;

        // 一个核总共需要处理的数 672
        uint32_t coreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;

        // 因为 单核一次计算 8 block，需要处理 42 个 block，共需要 42/8 向上取整 = 6 次
        uint32_t tileNum = everyCoreInputBlockNum / tileBlockNum;
        uint32_t finalTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? tileNum : tileNum + 1;

        // 尾块: 总数 - 前 tileNum 次处理的数，一次可以处理 tileDataNum 数，672 - 128 * 5
        uint32_t tailDataNum = coreDataNum - tileDataNum * tileNum;
        tailDataNum = tailDataNum == 0 ? tileDataNum : tailDataNum;

        tiling.set_coreDataNum(coreDataNum);   // 对齐后共需要处理的数
        tiling.set_tileDataNum(tileDataNum);   // 单核一次可以处理的数
        tiling.set_finalTileNum(finalTileNum); // 对齐后需要处理的次数
        tiling.set_tailDataNum(tailDataNum);   // 尾块需要处理的数

        context->SetBlockDim(1);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;

        return ge::GRAPH_SUCCESS;
    }
}

namespace ge
{
    static ge::graphStatus InferShape(gert::InferShapeContext *context)
    {
        const gert::Shape *x1_shape = context->GetInputShape(0);
        gert::Shape *y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;
        return GRAPH_SUCCESS;
    }
}

namespace ops
{
    class Asinh : public OpDef
    {
    public:
        explicit Asinh(const char *name) : OpDef(name)
        {
            this->Input("x")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

            this->SetInferShape(ge::InferShape);

            this->AICore()
                .SetTiling(optiling::TilingFunc);
            this->AICore().AddConfig("ascend910b").AddConfig("ascend310b");
        }
    };

    OP_ADD(Asinh);
}
